package main

import (
	"fmt"
)

/*
用例 14/17 超时
需要按照前缀某一前缀遍历Trie
*/
type WordFilter struct {
	preTrie *Trie
	sufTrie *Trie
}

func Constructor(words []string) WordFilter {
	preTrie := NewTrie()
	sufTrie := NewTrie()
	for i, word := range words {
		preTrie.addStr(word, i)
		sufTrie.addStr(reverse(word), i)
	}
	//fmt.Println(preTrie.containsStr("abbba"), preTrie.containsStr("abba"))
	//fmt.Println(sufTrie.containsStr("abbba"), sufTrie.containsStr("abba"))
	return WordFilter{
		preTrie: preTrie,
		sufTrie: sufTrie,
	}
}

func (this *WordFilter) F(pref string, suff string) int {
	map1 := this.preTrie.walkTree(pref)
	map2 := this.sufTrie.walkTree(reverse(suff))
	//fmt.Println(map1, map2)
	maxKey := retain(map1, map2)
	//fmt.Println(map1, map2, maxKey)
	return maxKey
}

func reverse(word string) string {
	N := len(word)
	res := make([]byte, N)
	for i, ch := range []byte(word) {
		res[N-1-i] = ch
	}
	return string(res)
}

/**
 * Your WordFilter object will be instantiated and called as such:
 * obj := Constructor(words);
 * param_1 := obj.F(pref,suff);
 */

type Trie struct {
	root *Node
}

type Node struct {
	children []*Node
	isWord   bool
	index    int //单词在原words数字中的下表
}

func NewTrie() *Trie {
	return &Trie{&Node{}}
}

func (t *Trie) addStr(str string, index int) {
	chs := []byte(str)
	node := t.root
	for _, ch := range chs {
		if node.children == nil {
			node.children = make([]*Node, 128)
		}
		if node.children[ch] == nil {
			node.children[ch] = &Node{}
		}
		node = node.children[ch]
	}
	node.isWord = true
	//记录原下标
	node.index = index
}

func (t *Trie) containsStr(str string) bool {
	chs := []byte(str)
	node := t.root
	for _, ch := range chs {
		if node.children == nil || node.children[ch] == nil {
			return false
		}
		node = node.children[ch]
	}
	return node.isWord
}

/*
该方法不符合题目要求
*/
func (t *Trie) getStrIndex(str string) int {
	chs := []byte(str)
	node := t.root
	for _, ch := range chs {
		if node.children == nil || node.children[ch] == nil {
			return -1
		}
		node = node.children[ch]
	}
	//return node.isWord
	if node.isWord {
		return node.index
	}
	return -1
}

/**
传入前缀pre，遍历Trie中所有符合前缀的单词，并返回这些单词的原下表
*/
func (t *Trie) walkTree(pre string) map[int]bool {
	chs := []byte(pre)
	node := t.root
	for _, ch := range chs {
		if node.children == nil || node.children[ch] == nil {
			return nil
		}
		node = node.children[ch]
	}
	//遍历node下面所有分支
	resMap := make(map[int]bool)
	dfs(node, resMap)
	return resMap
}

func dfs(root *Node, resMap map[int]bool) {
	if root == nil {
		return
	}
	if root.isWord {
		resMap[root.index] = true
	}
	for ch := 'a'; ch <= 'z'; ch++ {
		if root.children != nil && root.children[ch] != nil {
			dfs(root.children[ch], resMap)
		}
	}
}

/**
取两个map[int]bool的key的交集中最大的key，如果没有交集返回-1
*/
func retain(map1, map2 map[int]bool) int {
	maxKey := -1
	for k, _ := range map1 {
		_, contains := map2[k]
		if contains && k > maxKey {
			maxKey = k
		}
	}
	return maxKey
}

func main() {
	//trie := NewTrie()
	//trie.addStr("abc", 0)
	//trie.addStr("ab", 1)
	//trie.addStr("abcd", 2)
	//trie.addStr("bca", 3)
	//fmt.Println(trie.containsStr("abc"))
	//fmt.Println(trie.containsStr("ab"))
	//fmt.Println(trie.containsStr("bc"))
	//
	//l1 := trie.walkTree("a")
	//fmt.Println(l1)

	//map1 := make(map[int]bool)
	//map1[5] = true
	//map1[03] = true
	//map1[4] = true
	//map2 := make(map[int]bool)
	//map2[3] = true
	//map2[0] = true
	//fmt.Println(retain(map1, map2))

	filter := Constructor([]string{"abbba", "abba"})
	res := filter.F("ab", "ba")
	fmt.Println(res)
}
